test: refactor tests and fix inference edge cases
- Convert protocol and MoE test classes to plain functions - Add real server/engine integration and generate_async tests - Isolate test model per test and use pytest tmp_path - Reset FastAPI engine state after inference tests - Fix generate_async StopIteration handling on Python 3.12 - Fix HF adapter MoE dense/shared and Gemma qk_norm mapping - Correct dev dependency httpx2 to httpx
This commit is contained in:
+5
-9
@@ -1,7 +1,5 @@
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import tempfile
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -44,20 +42,18 @@ def test_tokenizer():
|
||||
return create_test_tokenizer()
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@pytest.fixture
|
||||
def test_model(device):
|
||||
"""Session-scoped small AutoRegressiveLM model, created once."""
|
||||
"""Function-scoped small AutoRegressiveLM model, isolated per test."""
|
||||
config = make_tiny_config()
|
||||
model = AutoRegressiveLM(config).to(device=device)
|
||||
return {"model": model, "device": device, "config": config}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_dir():
|
||||
"""Function-scoped temporary directory, cleaned up after each test."""
|
||||
d = tempfile.mkdtemp()
|
||||
yield d
|
||||
shutil.rmtree(d, ignore_errors=True)
|
||||
def temp_dir(tmp_path):
|
||||
"""Function-scoped temporary directory, cleaned up by pytest."""
|
||||
return str(tmp_path)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
@@ -8,6 +8,13 @@ from fastapi.testclient import TestClient
|
||||
from astrai.inference import get_app
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _cleanup_app_engine():
|
||||
"""Reset the lazy FastAPI singleton engine after each inference test."""
|
||||
yield
|
||||
get_app().state.engine = None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client():
|
||||
"""Provide a test client for the FastAPI app."""
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
"""Unit tests for GenerateResult accumulator and InferenceEngine.generate()."""
|
||||
|
||||
import asyncio
|
||||
import threading
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
@@ -156,6 +157,35 @@ def test_engine_generate_streaming_yields_tokens():
|
||||
assert tokens == ["t1", "t2"]
|
||||
|
||||
|
||||
def test_engine_generate_async_yields_tokens_until_stop():
|
||||
mock_model, mock_tokenizer = _make_engine_mocks(decode="tok")
|
||||
callbacks_saved = []
|
||||
|
||||
def capture_cb(prompt, **kw):
|
||||
callbacks_saved.append(kw.get("stream_callback"))
|
||||
|
||||
with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
|
||||
instance = MockSched.return_value
|
||||
instance.add_task.side_effect = capture_cb
|
||||
instance.remove_task.return_value = []
|
||||
|
||||
eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
|
||||
agen = eng.generate_async("hello")
|
||||
|
||||
async def collect():
|
||||
out = []
|
||||
async for token in agen:
|
||||
out.append(token)
|
||||
return out
|
||||
|
||||
cb = callbacks_saved[0]
|
||||
cb("t1")
|
||||
cb("t2")
|
||||
cb(STOP)
|
||||
|
||||
assert asyncio.run(collect()) == ["t1", "t2"]
|
||||
|
||||
|
||||
def test_engine_generate_non_streaming_batch():
|
||||
mock_model, mock_tokenizer = _make_engine_mocks(decode="r")
|
||||
|
||||
|
||||
+227
-205
@@ -3,8 +3,6 @@
|
||||
import json
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from astrai.inference.network.anthropic import AnthropicResponseBuilder
|
||||
from astrai.inference.network.openai import OpenAIResponseBuilder
|
||||
from astrai.inference.network.protocol import GenContext, StopChecker, StopInfo
|
||||
@@ -34,223 +32,247 @@ def _sse_payloads(events):
|
||||
return payloads
|
||||
|
||||
|
||||
class TestStopChecker:
|
||||
def test_check_finds_match(self):
|
||||
sc = StopChecker(["stop", "end"])
|
||||
assert sc.check("hello stop world") == "stop"
|
||||
|
||||
def test_check_returns_none_when_no_match(self):
|
||||
sc = StopChecker(["stop"])
|
||||
assert sc.check("hello world") is None
|
||||
|
||||
def test_check_empty_sequences(self):
|
||||
sc = StopChecker([])
|
||||
assert sc.check("hello") is None
|
||||
def _make_openai_builder():
|
||||
builder = OpenAIResponseBuilder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hello")]
|
||||
req.stop = None
|
||||
req.model = "astrai"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hello"
|
||||
builder.prepare(req, engine)
|
||||
return builder
|
||||
|
||||
|
||||
class TestGenContext:
|
||||
def test_defaults(self):
|
||||
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
|
||||
assert ctx.completion_tokens == 0
|
||||
|
||||
def test_fields_mutable(self):
|
||||
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
|
||||
ctx.completion_tokens = 42
|
||||
assert ctx.completion_tokens == 42
|
||||
def _make_anthropic_builder():
|
||||
builder = AnthropicResponseBuilder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hello")]
|
||||
req.model = "claude"
|
||||
req.system = None
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hello"
|
||||
builder.prepare(req, engine)
|
||||
return builder
|
||||
|
||||
|
||||
class TestStopInfo:
|
||||
def test_defaults(self):
|
||||
s = StopInfo()
|
||||
assert s.matched is None
|
||||
assert s.body == ""
|
||||
assert s.yielded == ""
|
||||
|
||||
def test_with_values(self):
|
||||
s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
|
||||
assert s.matched == "stop"
|
||||
assert s.body == "hello stop"
|
||||
assert s.yielded == "hello "
|
||||
def test_check_finds_match():
|
||||
sc = StopChecker(["stop", "end"])
|
||||
assert sc.check("hello stop world") == "stop"
|
||||
|
||||
|
||||
class TestOpenAIResponseBuilder:
|
||||
@pytest.fixture
|
||||
def builder(self):
|
||||
builder = OpenAIResponseBuilder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hello")]
|
||||
req.stop = None
|
||||
req.model = "astrai"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hello"
|
||||
builder.prepare(req, engine)
|
||||
return builder
|
||||
|
||||
def test_prepare_returns_prompt_ctx_stops(self, builder):
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hi")]
|
||||
req.stop = ["END"]
|
||||
req.model = "gpt"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hi"
|
||||
prompt, ctx, stops = builder.prepare(req, engine)
|
||||
assert prompt == "Hi"
|
||||
assert ctx.model == "gpt"
|
||||
assert ctx.prompt_tokens == 0
|
||||
assert stops == ["END"]
|
||||
|
||||
def test_prepare_no_stop_returns_empty_list(self, builder):
|
||||
req = MagicMock()
|
||||
req.messages = []
|
||||
req.stop = None
|
||||
req.model = "x"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = ""
|
||||
_, _, stops = builder.prepare(req, engine)
|
||||
assert stops == []
|
||||
|
||||
def test_format_stream_start(self, builder):
|
||||
ctx = _make_ctx()
|
||||
events = builder.format_stream_start(ctx)
|
||||
payloads = _sse_payloads(events)
|
||||
assert len(payloads) == 1
|
||||
p = payloads[0]
|
||||
assert p["object"] == "chat.completion.chunk"
|
||||
assert p["choices"][0]["delta"]["role"] == "assistant"
|
||||
assert p["choices"][0]["finish_reason"] is None
|
||||
|
||||
def test_format_chunk(self, builder):
|
||||
events = builder.format_chunk("hello", body="hello")
|
||||
payload = json.loads(events[0].split("data: ", 1)[1])
|
||||
assert payload["choices"][0]["delta"]["content"] == "hello"
|
||||
assert payload["choices"][0]["finish_reason"] is None
|
||||
|
||||
def test_format_stream_end(self, builder):
|
||||
ctx = _make_ctx(completion_tokens=5)
|
||||
stop = StopInfo(matched="stop")
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
finish = payloads[0]
|
||||
assert finish["choices"][0]["finish_reason"] == "stop"
|
||||
usage = payloads[1]
|
||||
assert usage["completion_tokens"] == 5
|
||||
assert usage["total_tokens"] == 15
|
||||
|
||||
def test_format_response(self, builder):
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo()
|
||||
resp = builder.format_response(ctx, "hello", stop)
|
||||
assert resp["object"] == "chat.completion"
|
||||
assert resp["choices"][0]["message"]["content"] == "hello"
|
||||
assert resp["usage"]["prompt_tokens"] == 10
|
||||
def test_check_returns_none_when_no_match():
|
||||
sc = StopChecker(["stop"])
|
||||
assert sc.check("hello world") is None
|
||||
|
||||
|
||||
class TestAnthropicResponseBuilder:
|
||||
@pytest.fixture
|
||||
def builder(self):
|
||||
builder = AnthropicResponseBuilder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hello")]
|
||||
req.model = "claude"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hello"
|
||||
req.system = None
|
||||
builder.prepare(req, engine)
|
||||
return builder
|
||||
def test_check_empty_sequences():
|
||||
sc = StopChecker([])
|
||||
assert sc.check("hello") is None
|
||||
|
||||
def test_prepare_messages(self, builder):
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hi")]
|
||||
req.model = "claude"
|
||||
req.system = None
|
||||
req.stop_sequences = None
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hi"
|
||||
prompt, ctx, stops = builder.prepare(req, engine)
|
||||
assert prompt == "Hi"
|
||||
assert stops == []
|
||||
|
||||
def test_prepare_with_stop_sequences(self, builder):
|
||||
req = MagicMock()
|
||||
req.messages = []
|
||||
req.model = "x"
|
||||
req.stop_sequences = ["stop", "end"]
|
||||
req.system = None
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = ""
|
||||
_, _, stops = builder.prepare(req, engine)
|
||||
assert stops == ["stop", "end"]
|
||||
def test_gen_context_defaults():
|
||||
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
|
||||
assert ctx.completion_tokens == 0
|
||||
|
||||
def test_format_stream_start(self, builder):
|
||||
ctx = _make_ctx(prompt_tokens=3)
|
||||
events = builder.format_stream_start(ctx)
|
||||
payloads = _sse_payloads(events)
|
||||
assert len(payloads) == 2
|
||||
assert payloads[0]["type"] == "message_start"
|
||||
assert payloads[0]["message"]["usage"]["input_tokens"] == 3
|
||||
assert payloads[1]["type"] == "content_block_start"
|
||||
|
||||
def test_format_chunk(self, builder):
|
||||
events = builder.format_chunk("tok", body="tok")
|
||||
payload = json.loads(events[0].split("data: ", 1)[1])
|
||||
assert payload["type"] == "content_block_delta"
|
||||
assert payload["delta"]["text"] == "tok"
|
||||
def test_gen_context_fields_mutable():
|
||||
ctx = GenContext(resp_id="a", created=1, model="m", prompt_tokens=10)
|
||||
ctx.completion_tokens = 42
|
||||
assert ctx.completion_tokens == 42
|
||||
|
||||
def test_format_stream_end_no_stop(self, builder):
|
||||
ctx = _make_ctx(completion_tokens=3)
|
||||
stop = StopInfo()
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
# content_block_stop, message_delta, message_stop
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == ["content_block_stop", "message_delta", "message_stop"]
|
||||
assert payloads[1]["delta"]["stop_reason"] == "end_turn"
|
||||
|
||||
def test_format_stream_end_with_stop_trims_and_emits_remaining(self, builder):
|
||||
ctx = _make_ctx(completion_tokens=7)
|
||||
stop = StopInfo(
|
||||
matched="END",
|
||||
body="Hello world END extra",
|
||||
yielded="Hello ",
|
||||
)
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
# unyielded delta, content_block_stop, message_delta, message_stop
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == [
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
assert payloads[0]["delta"]["text"] == "world "
|
||||
assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
|
||||
assert payloads[2]["delta"]["stop_sequence"] == "END"
|
||||
def test_stop_info_defaults():
|
||||
s = StopInfo()
|
||||
assert s.matched is None
|
||||
assert s.body == ""
|
||||
assert s.yielded == ""
|
||||
|
||||
def test_format_stream_end_stop_trimmed_already_yielded(self, builder):
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo(
|
||||
matched="END",
|
||||
body="Hello END",
|
||||
yielded="Hello ",
|
||||
)
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
# No unyielded delta (everything already sent)
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == ["content_block_stop", "message_delta", "message_stop"]
|
||||
|
||||
def test_format_response_with_stop_trims_content(self, builder):
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
|
||||
resp = builder.format_response(ctx, "text STOP extra", stop)
|
||||
assert resp["content"][0]["text"] == "text "
|
||||
assert resp["stop_reason"] == "stop_sequence"
|
||||
assert resp["stop_sequence"] == "STOP"
|
||||
def test_stop_info_with_values():
|
||||
s = StopInfo(matched="stop", body="hello stop", yielded="hello ")
|
||||
assert s.matched == "stop"
|
||||
assert s.body == "hello stop"
|
||||
assert s.yielded == "hello "
|
||||
|
||||
def test_format_response_no_stop(self, builder):
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo()
|
||||
resp = builder.format_response(ctx, "full text", stop)
|
||||
assert resp["content"][0]["text"] == "full text"
|
||||
assert resp["stop_reason"] == "end_turn"
|
||||
|
||||
def test_openai_prepare_returns_prompt_ctx_stops():
|
||||
builder = _make_openai_builder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hi")]
|
||||
req.stop = ["END"]
|
||||
req.model = "gpt"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hi"
|
||||
prompt, ctx, stops = builder.prepare(req, engine)
|
||||
assert prompt == "Hi"
|
||||
assert ctx.model == "gpt"
|
||||
assert ctx.prompt_tokens == 0
|
||||
assert stops == ["END"]
|
||||
|
||||
|
||||
def test_openai_prepare_no_stop_returns_empty_list():
|
||||
builder = _make_openai_builder()
|
||||
req = MagicMock()
|
||||
req.messages = []
|
||||
req.stop = None
|
||||
req.model = "x"
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = ""
|
||||
_, _, stops = builder.prepare(req, engine)
|
||||
assert stops == []
|
||||
|
||||
|
||||
def test_openai_format_stream_start():
|
||||
builder = _make_openai_builder()
|
||||
ctx = _make_ctx()
|
||||
events = builder.format_stream_start(ctx)
|
||||
payloads = _sse_payloads(events)
|
||||
assert len(payloads) == 1
|
||||
p = payloads[0]
|
||||
assert p["object"] == "chat.completion.chunk"
|
||||
assert p["choices"][0]["delta"]["role"] == "assistant"
|
||||
assert p["choices"][0]["finish_reason"] is None
|
||||
|
||||
|
||||
def test_openai_format_chunk():
|
||||
builder = _make_openai_builder()
|
||||
events = builder.format_chunk("hello", body="hello")
|
||||
payload = json.loads(events[0].split("data: ", 1)[1])
|
||||
assert payload["choices"][0]["delta"]["content"] == "hello"
|
||||
assert payload["choices"][0]["finish_reason"] is None
|
||||
|
||||
|
||||
def test_openai_format_stream_end():
|
||||
builder = _make_openai_builder()
|
||||
ctx = _make_ctx(completion_tokens=5)
|
||||
stop = StopInfo(matched="stop")
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
finish = payloads[0]
|
||||
assert finish["choices"][0]["finish_reason"] == "stop"
|
||||
usage = payloads[1]
|
||||
assert usage["completion_tokens"] == 5
|
||||
assert usage["total_tokens"] == 15
|
||||
|
||||
|
||||
def test_openai_format_response():
|
||||
builder = _make_openai_builder()
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo()
|
||||
resp = builder.format_response(ctx, "hello", stop)
|
||||
assert resp["object"] == "chat.completion"
|
||||
assert resp["choices"][0]["message"]["content"] == "hello"
|
||||
assert resp["usage"]["prompt_tokens"] == 10
|
||||
|
||||
|
||||
def test_anthropic_prepare_messages():
|
||||
builder = _make_anthropic_builder()
|
||||
req = MagicMock()
|
||||
req.messages = [MagicMock(role="user", content="Hi")]
|
||||
req.model = "claude"
|
||||
req.system = None
|
||||
req.stop_sequences = None
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = "Hi"
|
||||
prompt, ctx, stops = builder.prepare(req, engine)
|
||||
assert prompt == "Hi"
|
||||
assert stops == []
|
||||
|
||||
|
||||
def test_anthropic_prepare_with_stop_sequences():
|
||||
builder = _make_anthropic_builder()
|
||||
req = MagicMock()
|
||||
req.messages = []
|
||||
req.model = "x"
|
||||
req.stop_sequences = ["stop", "end"]
|
||||
req.system = None
|
||||
engine = MagicMock()
|
||||
engine.tokenizer.apply_chat_template.return_value = ""
|
||||
_, _, stops = builder.prepare(req, engine)
|
||||
assert stops == ["stop", "end"]
|
||||
|
||||
|
||||
def test_anthropic_format_stream_start():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx(prompt_tokens=3)
|
||||
events = builder.format_stream_start(ctx)
|
||||
payloads = _sse_payloads(events)
|
||||
assert len(payloads) == 2
|
||||
assert payloads[0]["type"] == "message_start"
|
||||
assert payloads[0]["message"]["usage"]["input_tokens"] == 3
|
||||
assert payloads[1]["type"] == "content_block_start"
|
||||
|
||||
|
||||
def test_anthropic_format_chunk():
|
||||
builder = _make_anthropic_builder()
|
||||
events = builder.format_chunk("tok", body="tok")
|
||||
payload = json.loads(events[0].split("data: ", 1)[1])
|
||||
assert payload["type"] == "content_block_delta"
|
||||
assert payload["delta"]["text"] == "tok"
|
||||
|
||||
|
||||
def test_anthropic_format_stream_end_no_stop():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx(completion_tokens=3)
|
||||
stop = StopInfo()
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == ["content_block_stop", "message_delta", "message_stop"]
|
||||
assert payloads[1]["delta"]["stop_reason"] == "end_turn"
|
||||
|
||||
|
||||
def test_anthropic_format_stream_end_with_stop_trims_and_emits_remaining():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx(completion_tokens=7)
|
||||
stop = StopInfo(
|
||||
matched="END",
|
||||
body="Hello world END extra",
|
||||
yielded="Hello ",
|
||||
)
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == [
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
assert payloads[0]["delta"]["text"] == "world "
|
||||
assert payloads[2]["delta"]["stop_reason"] == "stop_sequence"
|
||||
assert payloads[2]["delta"]["stop_sequence"] == "END"
|
||||
|
||||
|
||||
def test_anthropic_format_stream_end_stop_trimmed_already_yielded():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo(
|
||||
matched="END",
|
||||
body="Hello END",
|
||||
yielded="Hello ",
|
||||
)
|
||||
events = builder.format_stream_end(ctx, stop)
|
||||
payloads = _sse_payloads(events)
|
||||
types = [p["type"] for p in payloads]
|
||||
assert types == ["content_block_stop", "message_delta", "message_stop"]
|
||||
|
||||
|
||||
def test_anthropic_format_response_with_stop_trims_content():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo(matched="STOP", body="text STOP extra", yielded="text ")
|
||||
resp = builder.format_response(ctx, "text STOP extra", stop)
|
||||
assert resp["content"][0]["text"] == "text "
|
||||
assert resp["stop_reason"] == "stop_sequence"
|
||||
assert resp["stop_sequence"] == "STOP"
|
||||
|
||||
|
||||
def test_anthropic_format_response_no_stop():
|
||||
builder = _make_anthropic_builder()
|
||||
ctx = _make_ctx()
|
||||
stop = StopInfo()
|
||||
resp = builder.format_response(ctx, "full text", stop)
|
||||
assert resp["content"][0]["text"] == "full text"
|
||||
assert resp["stop_reason"] == "end_turn"
|
||||
|
||||
@@ -1,8 +1,15 @@
|
||||
"""Unit tests for the inference HTTP server."""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from astrai.inference import get_app
|
||||
from astrai.inference.network.app import _create_engine
|
||||
from astrai.model.transformer import AutoRegressiveLM
|
||||
from astrai.serialization import save_model
|
||||
from tests.helpers import CHAT_TEMPLATE, build_test_tokenizer, make_tiny_config
|
||||
|
||||
|
||||
def test_health_no_model(client):
|
||||
@@ -212,5 +219,42 @@ def test_chat_completions_stop_sequence_stream(client, loaded_model):
|
||||
)
|
||||
|
||||
|
||||
def test_chat_completions_real_engine(tmp_path, client):
|
||||
"""POST /v1/chat/completions with a real tiny model and tokenizer."""
|
||||
cfg = make_tiny_config(vocab_size=256)
|
||||
model = AutoRegressiveLM(cfg).eval()
|
||||
save_model(cfg.to_dict(), model.state_dict(), str(tmp_path))
|
||||
|
||||
tokenizer = build_test_tokenizer(vocab_size=256, chat_template=CHAT_TEMPLATE)
|
||||
tokenizer.save_pretrained(str(tmp_path))
|
||||
|
||||
engine = _create_engine(
|
||||
Path(tmp_path),
|
||||
device="cpu",
|
||||
dtype=torch.float32,
|
||||
max_batch_size=1,
|
||||
max_seq_len=64,
|
||||
)
|
||||
try:
|
||||
get_app().state.engine = engine
|
||||
response = client.post(
|
||||
"/v1/chat/completions",
|
||||
json={
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"max_tokens": 4,
|
||||
"temperature": 0.0,
|
||||
"stream": False,
|
||||
},
|
||||
)
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
content = data["choices"][0]["message"]["content"]
|
||||
assert isinstance(content, str)
|
||||
assert data["usage"]["completion_tokens"] > 0
|
||||
finally:
|
||||
engine.shutdown()
|
||||
get_app().state.engine = None
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
+169
-161
@@ -54,6 +54,20 @@ def _make_batch(config, batch_size=2, seq_len=8, with_extra=False):
|
||||
return batch
|
||||
|
||||
|
||||
def _make_seq_moe_fixture(device):
|
||||
config = _make_tiny_moe_config()
|
||||
model = _make_model(config).to(device)
|
||||
model.train()
|
||||
return config, model
|
||||
|
||||
|
||||
def _make_sft_moe_fixture(device):
|
||||
config = _make_tiny_moe_config()
|
||||
model = _make_model(config).to(device)
|
||||
model.train()
|
||||
return config, model
|
||||
|
||||
|
||||
def test_model_forward_contract_uses_dense_training_and_packed_inference():
|
||||
from astrai.inference.cache import PagePool, TaskCacheManager
|
||||
from astrai.inference.workspace import InferenceWorkspace
|
||||
@@ -185,172 +199,166 @@ def test_moe_metrics_flow_through_wrapped_model(device):
|
||||
assert "router_entropy" in strategy._moe_metrics
|
||||
|
||||
|
||||
class TestSEQStrategyMoE:
|
||||
"""End‑to‑end tests for SEQStrategy with MoE aux loss."""
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, device):
|
||||
self.device = device
|
||||
self.config = _make_tiny_moe_config()
|
||||
self.model = _make_model(self.config).to(device)
|
||||
self.model.train()
|
||||
|
||||
def test_compute_loss_returns_scalar(self):
|
||||
"""compute_loss should return a scalar tensor."""
|
||||
strategy = SEQStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
loss = strategy.compute_loss(_make_batch(self.config))
|
||||
assert loss.ndim == 0
|
||||
assert loss.requires_grad
|
||||
|
||||
def test_compute_loss_output_has_metrics(self):
|
||||
"""compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
|
||||
strategy = SEQStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config))
|
||||
|
||||
assert "loss" in output
|
||||
assert "metrics" in output
|
||||
assert output["loss"].ndim == 0
|
||||
assert output["loss"].requires_grad
|
||||
|
||||
metrics = output["metrics"]
|
||||
# MoE metrics should appear when coef > 0 and model has MoE layers
|
||||
for key in ("moe_aux_loss", "moe_aux_loss_weighted", "task_loss", "loss"):
|
||||
assert key in metrics, f"Missing metric: {key}"
|
||||
assert isinstance(metrics[key], float)
|
||||
|
||||
def test_moe_metrics_populated_after_forward(self):
|
||||
"""strategy._moe_metrics populated after compute_loss_output."""
|
||||
strategy = SEQStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
strategy.compute_loss_output(_make_batch(self.config))
|
||||
|
||||
moe_metrics = strategy._moe_metrics
|
||||
assert moe_metrics, "_moe_metrics should not be empty for MoE model"
|
||||
for key in (
|
||||
"aux_loss",
|
||||
"router_entropy",
|
||||
"dead_expert_fraction",
|
||||
"load_imbalance_mean",
|
||||
"load_imbalance_max",
|
||||
):
|
||||
assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
|
||||
assert isinstance(moe_metrics[key], float)
|
||||
|
||||
def test_zero_coef_zeroes_weighted_aux(self):
|
||||
"""moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
|
||||
strategy = SEQStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.0,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config))
|
||||
metrics = output["metrics"]
|
||||
|
||||
# task_loss and loss should be equal (aux weighted by zero)
|
||||
assert "task_loss" in metrics
|
||||
assert "loss" in metrics
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
|
||||
# weighted aux loss is zero
|
||||
assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
|
||||
|
||||
# MoE diagnostics are still collected (monitoring purposes)
|
||||
assert strategy._moe_metrics
|
||||
assert "router_entropy" in strategy._moe_metrics
|
||||
|
||||
def test_aux_loss_added_to_total_loss(self):
|
||||
"""Total loss > task_loss when moe_aux_loss_coef > 0."""
|
||||
strategy = SEQStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config))
|
||||
assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
|
||||
|
||||
def test_factory_creates_strategy_with_coef(self):
|
||||
"""StrategyFactory.create passes moe_aux_loss_coef to strategy."""
|
||||
strategy = StrategyFactory.create(
|
||||
"seq",
|
||||
model=self.model,
|
||||
device=self.device,
|
||||
moe_aux_loss_coef=0.02,
|
||||
)
|
||||
assert strategy.moe_aux_loss_coef == 0.02
|
||||
|
||||
def test_no_aux_loss_for_mlp_model(self):
|
||||
"""Pure MLP model: model outputs no aux_loss → no MoE metrics."""
|
||||
from astrai.config.model_config import AutoRegressiveLMConfig
|
||||
|
||||
mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
|
||||
mlp_model = AutoRegressiveLM(mlp_config).to(self.device)
|
||||
mlp_model.train()
|
||||
|
||||
strategy = SEQStrategy(
|
||||
mlp_model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config))
|
||||
metrics = output["metrics"]
|
||||
|
||||
assert "moe_aux_loss" not in metrics
|
||||
assert "moe_aux_loss_weighted" not in metrics
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
assert strategy._moe_metrics == {}
|
||||
def test_seq_compute_loss_returns_scalar(device):
|
||||
"""compute_loss should return a scalar tensor."""
|
||||
config, model = _make_seq_moe_fixture(device)
|
||||
strategy = SEQStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
loss = strategy.compute_loss(_make_batch(config))
|
||||
assert loss.ndim == 0
|
||||
assert loss.requires_grad
|
||||
|
||||
|
||||
class TestSFTStrategyMoE:
|
||||
"""End‑to‑end tests for SFTStrategy with MoE aux loss."""
|
||||
def test_seq_compute_loss_output_has_metrics(device):
|
||||
"""compute_loss_output dict with moe_aux_loss_coef > 0 includes MoE metrics."""
|
||||
config, model = _make_seq_moe_fixture(device)
|
||||
strategy = SEQStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config))
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup(self, device):
|
||||
self.device = device
|
||||
self.config = _make_tiny_moe_config()
|
||||
self.model = _make_model(self.config).to(device)
|
||||
self.model.train()
|
||||
assert "loss" in output
|
||||
assert "metrics" in output
|
||||
assert output["loss"].ndim == 0
|
||||
assert output["loss"].requires_grad
|
||||
|
||||
def test_compute_loss_output_with_aux_loss(self):
|
||||
"""SFTStrategy produces MoE metrics when coef > 0."""
|
||||
strategy = SFTStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
|
||||
metrics = output["metrics"]
|
||||
# MoE metrics should appear when coef > 0 and model has MoE layers
|
||||
for key in ("moe_aux_loss", "moe_aux_loss_weighted", "task_loss", "loss"):
|
||||
assert key in metrics, f"Missing metric: {key}"
|
||||
assert isinstance(metrics[key], float)
|
||||
|
||||
metrics = output["metrics"]
|
||||
assert "moe_aux_loss" in metrics
|
||||
assert "moe_aux_loss_weighted" in metrics
|
||||
assert metrics["loss"] > metrics["task_loss"] + 1e-12
|
||||
|
||||
moe_metrics = strategy._moe_metrics
|
||||
assert "router_entropy" in moe_metrics
|
||||
assert "dead_expert_fraction" in moe_metrics
|
||||
def test_seq_moe_metrics_populated_after_forward(device):
|
||||
"""strategy._moe_metrics populated after compute_loss_output."""
|
||||
config, model = _make_seq_moe_fixture(device)
|
||||
strategy = SEQStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
strategy.compute_loss_output(_make_batch(config))
|
||||
|
||||
def test_sft_zero_coef_zeroes_weighted_aux(self):
|
||||
"""SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
|
||||
strategy = SFTStrategy(
|
||||
self.model,
|
||||
self.device,
|
||||
moe_aux_loss_coef=0.0,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(self.config, with_extra=True))
|
||||
metrics = output["metrics"]
|
||||
moe_metrics = strategy._moe_metrics
|
||||
assert moe_metrics, "_moe_metrics should not be empty for MoE model"
|
||||
for key in (
|
||||
"aux_loss",
|
||||
"router_entropy",
|
||||
"dead_expert_fraction",
|
||||
"load_imbalance_mean",
|
||||
"load_imbalance_max",
|
||||
):
|
||||
assert key in moe_metrics, f"Missing _moe_metrics key: {key}"
|
||||
assert isinstance(moe_metrics[key], float)
|
||||
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
|
||||
# Diagnostics still collected
|
||||
assert strategy._moe_metrics
|
||||
assert "router_entropy" in strategy._moe_metrics
|
||||
|
||||
def test_seq_zero_coef_zeroes_weighted_aux(device):
|
||||
"""moe_aux_loss_coef=0 → weighted_aux_loss is zero, task_loss == loss."""
|
||||
config, model = _make_seq_moe_fixture(device)
|
||||
strategy = SEQStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.0,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config))
|
||||
metrics = output["metrics"]
|
||||
|
||||
# task_loss and loss should be equal (aux weighted by zero)
|
||||
assert "task_loss" in metrics
|
||||
assert "loss" in metrics
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
|
||||
# weighted aux loss is zero
|
||||
assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
|
||||
|
||||
# MoE diagnostics are still collected (monitoring purposes)
|
||||
assert strategy._moe_metrics
|
||||
assert "router_entropy" in strategy._moe_metrics
|
||||
|
||||
|
||||
def test_seq_aux_loss_added_to_total_loss(device):
|
||||
"""Total loss > task_loss when moe_aux_loss_coef > 0."""
|
||||
config, model = _make_seq_moe_fixture(device)
|
||||
strategy = SEQStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config))
|
||||
assert output["metrics"]["loss"] > output["metrics"]["task_loss"] + 1e-12
|
||||
|
||||
|
||||
def test_seq_factory_creates_strategy_with_coef(device):
|
||||
"""StrategyFactory.create passes moe_aux_loss_coef to strategy."""
|
||||
_, model = _make_seq_moe_fixture(device)
|
||||
strategy = StrategyFactory.create(
|
||||
"seq",
|
||||
model=model,
|
||||
device=device,
|
||||
moe_aux_loss_coef=0.02,
|
||||
)
|
||||
assert strategy.moe_aux_loss_coef == 0.02
|
||||
|
||||
|
||||
def test_seq_no_aux_loss_for_mlp_model(device):
|
||||
"""Pure MLP model: model outputs no aux_loss → no MoE metrics."""
|
||||
config, _ = _make_seq_moe_fixture(device)
|
||||
mlp_config = AutoRegressiveLMConfig(**{**TINY_CONFIG, "ffn_type": "mlp"})
|
||||
mlp_model = AutoRegressiveLM(mlp_config).to(device)
|
||||
mlp_model.train()
|
||||
|
||||
strategy = SEQStrategy(
|
||||
mlp_model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config))
|
||||
metrics = output["metrics"]
|
||||
|
||||
assert "moe_aux_loss" not in metrics
|
||||
assert "moe_aux_loss_weighted" not in metrics
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
assert strategy._moe_metrics == {}
|
||||
|
||||
|
||||
def test_sft_compute_loss_output_with_aux_loss(device):
|
||||
"""SFTStrategy produces MoE metrics when coef > 0."""
|
||||
config, model = _make_sft_moe_fixture(device)
|
||||
strategy = SFTStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.01,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
|
||||
|
||||
metrics = output["metrics"]
|
||||
assert "moe_aux_loss" in metrics
|
||||
assert "moe_aux_loss_weighted" in metrics
|
||||
assert metrics["loss"] > metrics["task_loss"] + 1e-12
|
||||
|
||||
moe_metrics = strategy._moe_metrics
|
||||
assert "router_entropy" in moe_metrics
|
||||
assert "dead_expert_fraction" in moe_metrics
|
||||
|
||||
|
||||
def test_sft_zero_coef_zeroes_weighted_aux(device):
|
||||
"""SFTStrategy with zero coef: weighted aux is zero, loss == task_loss."""
|
||||
config, model = _make_sft_moe_fixture(device)
|
||||
strategy = SFTStrategy(
|
||||
model,
|
||||
device,
|
||||
moe_aux_loss_coef=0.0,
|
||||
)
|
||||
output = strategy.compute_loss_output(_make_batch(config, with_extra=True))
|
||||
metrics = output["metrics"]
|
||||
|
||||
assert metrics["loss"] == pytest.approx(metrics["task_loss"], abs=1e-6)
|
||||
assert metrics.get("moe_aux_loss_weighted") == pytest.approx(0.0, abs=1e-6)
|
||||
# Diagnostics still collected
|
||||
assert strategy._moe_metrics
|
||||
assert "router_entropy" in strategy._moe_metrics
|
||||
|
||||
@@ -183,6 +183,68 @@ def test_convert_hf_weights_rejects_mla():
|
||||
convert_hf_weights(sd, cfg)
|
||||
|
||||
|
||||
def test_convert_hf_config_qwen2_moe_preserves_sparse_fields():
|
||||
raw = {
|
||||
**LLAMA_RAW,
|
||||
"model_type": "qwen2_moe",
|
||||
"num_local_experts": 2,
|
||||
"num_experts_per_tok": 1,
|
||||
"n_shared_experts": 1,
|
||||
"decoder_sparse_step": 2,
|
||||
"mlp_only_layers": [0],
|
||||
}
|
||||
cfg = ConfigFactory.load(convert_hf_config(raw))
|
||||
assert cfg.decoder_sparse_step == 2
|
||||
assert cfg.mlp_only_layers == [0]
|
||||
|
||||
|
||||
def test_convert_hf_config_gemma_enables_qk_norm():
|
||||
raw = {**LLAMA_RAW, "model_type": "gemma"}
|
||||
cfg = ConfigFactory.load(convert_hf_config(raw))
|
||||
assert cfg.use_qk_norm is True
|
||||
|
||||
|
||||
def test_convert_hf_weights_moe_with_dense_layers_roundtrip():
|
||||
cfg = make_tiny_config(
|
||||
ffn_type="moe",
|
||||
n_routed_experts=2,
|
||||
n_shared_experts=1,
|
||||
n_activated_experts=1,
|
||||
moe_intermediate_size=16,
|
||||
shared_expert_intermediate_size=16,
|
||||
mlp_only_layers=[0],
|
||||
decoder_sparse_step=1,
|
||||
)
|
||||
model = AutoRegressiveLM(cfg)
|
||||
converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
|
||||
assert_state_dicts_equal(converted, model.state_dict())
|
||||
|
||||
|
||||
def test_convert_hf_weights_qwen2_moe_singular_shared_expert_roundtrip():
|
||||
cfg = make_tiny_config(
|
||||
ffn_type="moe",
|
||||
n_routed_experts=2,
|
||||
n_shared_experts=1,
|
||||
n_activated_experts=1,
|
||||
moe_intermediate_size=16,
|
||||
shared_expert_intermediate_size=16,
|
||||
)
|
||||
model = AutoRegressiveLM(cfg)
|
||||
hf_sd = to_hf_keys(model.state_dict())
|
||||
hf_sd = {
|
||||
k.replace("shared_experts.", "shared_expert.", 1): v for k, v in hf_sd.items()
|
||||
}
|
||||
converted = convert_hf_weights(hf_sd, cfg)
|
||||
assert_state_dicts_equal(converted, model.state_dict())
|
||||
|
||||
|
||||
def test_convert_hf_weights_gemma_qk_norm_roundtrip():
|
||||
cfg = make_tiny_config(use_qk_norm=True)
|
||||
model = AutoRegressiveLM(cfg)
|
||||
converted = convert_hf_weights(to_hf_keys(model.state_dict()), cfg)
|
||||
assert_state_dicts_equal(converted, model.state_dict())
|
||||
|
||||
|
||||
def test_from_pretrained_hf_directory(tmp_path):
|
||||
cfg = make_tiny_config()
|
||||
model = AutoRegressiveLM(cfg).eval()
|
||||
|
||||
@@ -31,13 +31,11 @@ def test_gradient_checkpointing_empty_modules_noop(test_model):
|
||||
model = test_model["model"]
|
||||
callback = GradientCheckpointingCallback()
|
||||
|
||||
originals = [layer.forward for layer in model.layers]
|
||||
|
||||
for layer in model.layers:
|
||||
callback._enable(layer)
|
||||
|
||||
for layer, orig in zip(model.layers, originals):
|
||||
assert layer.forward is orig
|
||||
for layer in model.layers:
|
||||
assert not hasattr(layer, "_original_forward")
|
||||
|
||||
|
||||
def test_gradient_checkpointing_forward_unchanged(test_model):
|
||||
|
||||
Reference in New Issue
Block a user